import numpy as np
import sklearn.tree as tree
import matplotlib.pyplot as plt


def load_data(file):
    """
    从数据集中读取文件并转化为X矩阵和Y向量
    :param file: 要被读取的数据集
    :return: 返回X矩阵和Y向量
    """
    with open(file) as f:
        data = np.loadtxt(f, dtype=float, comments='#', delimiter=',')
        len_ = data.shape[1]
        X = data[:, :len_ - 1]
        # 数据集最后一列为Y并转化为向量
        Y = data[:, len_ - 1].reshape(-1, 1)
        return X, Y


def Tree_clf(x_train, y_train, x_test, y_test, criterion, max_depth):
    """
    创建决策树，生成训练集上决策树并输出精度Accuracy，绘制并保存决策树图像
    :param x_train: 训练集参数
    :param y_train: 训练集lable
    :param x_test: 验证集参数
    :param y_test: 验证集lable
    :param criterion: 划分选择标准
    :param max_depth:决策树最大层数
    :return:无返回，输出Accuracy并绘制图形
    """
    # 生成决策树
    clf = tree.DecisionTreeClassifier(random_state=1, criterion=criterion, max_depth=max_depth)
    clf = clf.fit(x_train, y_train)
    score = clf.score(x_test, y_test)
    print(f'Score in {criterion},{max_depth} is :{score:.5f}')
    # 绘制图形
    plt.figure()
    tree.plot_tree(decision_tree=clf, filled=True)
    plt.savefig(f'img/{criterion}_{max_depth}.png')

